import os
import sys
import subprocess
import shutil
import argparse
import shlex
from pathlib import Path

#  python3 /home/ubuntu/llmafl/generate_seeds.py --server-cmd "/home/ubuntu/experiments/wolfssl/build/examples/server/server --root /home/ubuntu/experiments/wolfssl -p 4433 -x -g -i -r -v 3 --fix-random -k /home/ubuntu/experiments/certs/server-key.pem -c /home/ubuntu/experiments/certs/server-cert.pem -d" --save-dir /home/ubuntu/experiments/out-wolfssl-pyafl/tlsfuzzer/tls12-oneway/save --seeds-dir /home/ubuntu/experiments/out-wolfssl-pyafl/tlsfuzzer/tls12-oneway/seeds --host 127.0.0.1 --port 4433 


PROJECT_ROOT = "/home/ubuntu/llmafl"
TLSFUZZER_DIR = os.path.join(PROJECT_ROOT, "tlsfuzzer")
TOOLS_DIR = os.path.join(TLSFUZZER_DIR, "tools")
SCRIPTS_DIR = os.path.join(TLSFUZZER_DIR, "scripts")
SEEDS_DIR = os.path.join(PROJECT_ROOT, "seeds")
VENV_DIR = os.path.join(TLSFUZZER_DIR, ".venv")

def run_cmd(cmd, cwd=None):
    print(f"Running: {cmd}")
    result = subprocess.run(cmd, shell=True, cwd=cwd, check=True)
    return result

def run_cmd_args(args, cwd=None):
    printable = " ".join(shlex.quote(str(a)) for a in args)
    print(f"Running: {printable}")
    result = subprocess.run(args, cwd=cwd, check=True)
    return result

def ensure_virtualenv():
    try:
        import virtualenv
    except ImportError:
        run_cmd("pip install virtualenv")

    if not os.path.exists(VENV_DIR):
        run_cmd(f"virtualenv .venv", cwd=TLSFUZZER_DIR)
    python_bin = os.path.join(VENV_DIR, "bin", "python")
    pip_bin = os.path.join(VENV_DIR, "bin", "pip")
    run_cmd(f"{pip_bin} install -r requirements.txt", cwd=TLSFUZZER_DIR)
    return python_bin

def build_arg_parser():
    p = argparse.ArgumentParser(description="Run tlsfuzzer and extract seeds without JSON config.")
    p.add_argument("--server-cmd", required=True, help="OpenSSL s_server command (quoted)")
    p.add_argument("--save-dir", required=True, help="Directory to save transcripts and logs")
    p.add_argument("--seeds-dir", default=SEEDS_DIR, help=f"Directory to place deduplicated seeds (default: {SEEDS_DIR})")
    p.add_argument("--scripts-dir", default=SCRIPTS_DIR, help=f"tlsfuzzer scripts directory (default: {SCRIPTS_DIR})")
    p.add_argument("--host", default="127.0.0.1", help="Server host (default: 127.0.0.1)")
    p.add_argument("--port", type=int, default=4433, help="Server port (default: 4433)")
    return p

def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    server_cmd = args.server_cmd
    save_output_root = args.save_dir
    scripts_dir = args.scripts_dir
    seeds_dir = args.seeds_dir
    host = args.host
    port = args.port

    ensure_virtualenv()
    python_bin = os.path.join(VENV_DIR, "bin", "python")

    # 2. 运行 run_all.py
    run_all_args = [
        python_bin,
        "run_all.py",
        "--save-dir", save_output_root,
        "--server-cmd", server_cmd,
        "--scripts-dir", scripts_dir,
        "--host", host,
        "--port", str(port),
    ]
    run_cmd_args(run_all_args, cwd=TOOLS_DIR)

    # 3. 运行 extract_send_bins.py，直接对 save_dir 递归收集 send.bin
    os.makedirs(seeds_dir, exist_ok=True)
    extract_cmd = [python_bin, "extract_send_bins.py", save_output_root, seeds_dir]
    run_cmd_args(extract_cmd, cwd=TOOLS_DIR)

    print(f"种子提取完毕，全部保存在: {seeds_dir}")

if __name__ == "__main__":
    main()